#include "compi/compi.h"

#include <gtest/gtest.h>

#include <mpi.h>
#include <string>
#include <vector>

using namespace compi;
using namespace compi::collectives;

// 全局测试环境指针，由自定义 main 显式创建与销毁，确保每个进程都调用 MPI_Init/Finalize。
static Environment * g_test_env = nullptr;

class CollectivesTest : public ::testing::Test
{
protected:
    void SetUp() override
    {
        // 借用全局环境，确保所有进程都调用 MPI_Init，并在套件结束时统一 MPI_Finalize。
        env_ = g_test_env;
        comm_ = &Context::for_comm(MPI_COMM_WORLD, *env_);
    }

    void TearDown() override
    {
        // 非拥有的指针清空；保留全局 Environment 直到进程退出。
        comm_ = nullptr;
        env_ = nullptr;
    }

    Environment * env_{};
    Context * comm_{}; // 非拥有的上下文指针（单例）
};

// 自定义 main：显式管理 MPI 生命周期，避免 mpiexec 认为未正确 finalize。
int
main(int argc, char ** argv)
{
    ::testing::InitGoogleTest(&argc, argv);
    // 构造全局环境，触发 MPI_Init_thread（若尚未初始化）
    g_test_env = new Environment();
    const int code = RUN_ALL_TESTS();
    // 显式析构环境，触发 MPI_Finalize
    delete g_test_env;
    g_test_env = nullptr;
    return code;
}

// broadcast: 固定类型
TEST_F(CollectivesTest, BroadcastFixedInt)
{
    int root = 0;
    int rank = comm_->rank();
    int value = (rank == root) ? 42 : -1;
    broadcast(*comm_, root, value);
    EXPECT_EQ(value, 42);
}

// broadcast: 非固定类型（字符串）
TEST_F(CollectivesTest, BroadcastString)
{
    int root = 0;
    int rank = comm_->rank();
    std::string payload = (rank == root) ? std::string("hello_collectives") : std::string();
    broadcast(*comm_, root, payload);
    EXPECT_EQ(payload, std::string("hello_collectives"));
}

// gather: 固定类型
TEST_F(CollectivesTest, GatherFixedInt)
{
    int root = 0;
    int rank = comm_->rank();
    int size = comm_->size();
    auto result = gather(*comm_, root, rank);
    if (rank == root)
    {
        ASSERT_EQ(static_cast<int>(result.size()), size);
        for (int i = 0; i < size; ++i)
            EXPECT_EQ(result[i], i);
    }
    else
    {
        EXPECT_TRUE(result.empty());
    }
}

// gather: 非固定类型（字符串）
TEST_F(CollectivesTest, GatherString)
{
    int root = 0;
    int rank = comm_->rank();
    int size = comm_->size();
    std::string s = std::string("rank_") + std::to_string(rank);
    auto result = gather(*comm_, root, s);
    if (rank == root)
    {
        ASSERT_EQ(static_cast<int>(result.size()), size);
        for (int i = 0; i < size; ++i)
        {
            EXPECT_EQ(result[i], std::string("rank_") + std::to_string(i));
        }
    }
    else
    {
        EXPECT_TRUE(result.empty());
    }
}

// alltoall: 固定类型（默认同步策略）
TEST_F(CollectivesTest, AlltoallFixedDefault)
{
    int rank = comm_->rank();
    int size = comm_->size();
    std::vector<int> send_values(size);
    for (int j = 0; j < size; ++j)
        send_values[j] = rank * 100 + j; // 发送到 j 的值

    auto recv_values = alltoall<int>(*comm_, send_values, SyncType::Default, SendMode::Default);
    ASSERT_EQ(static_cast<int>(recv_values.size()), size);
    for (int i = 0; i < size; ++i)
    {
        // 从 i 接收的值应该是 i*100 + rank
        EXPECT_EQ(recv_values[i], i * 100 + rank);
    }
}

// alltoall: 固定类型（NBX + 同步发送）
TEST_F(CollectivesTest, AlltoallFixedNBXSync)
{
    int rank = comm_->rank();
    int size = comm_->size();
    std::vector<int> send_values(size);
    for (int j = 0; j < size; ++j)
        send_values[j] = rank * 100 + j;

    auto recv_values = alltoall<int>(*comm_, send_values, SyncType::NBX, SendMode::Synchronous);
    ASSERT_EQ(static_cast<int>(recv_values.size()), size);
    for (int i = 0; i < size; ++i)
        EXPECT_EQ(recv_values[i], i * 100 + rank);
}

// alltoall: 非固定类型（字符串，默认策略 counts+Alltoallv）
TEST_F(CollectivesTest, AlltoallStringDefault)
{
    int rank = comm_->rank();
    int size = comm_->size();
    std::vector<std::string> send_values(size);
    for (int j = 0; j < size; ++j)
        send_values[j] = std::string("from_") + std::to_string(rank) + "_to_" + std::to_string(j);

    auto recv_values =
        alltoall<std::string>(*comm_, send_values, SyncType::Default, SendMode::Default);
    ASSERT_EQ(static_cast<int>(recv_values.size()), size);
    for (int i = 0; i < size; ++i)
    {
        std::string expected =
            std::string("from_") + std::to_string(i) + "_to_" + std::to_string(rank);
        EXPECT_EQ(recv_values[i], expected);
    }
}

// alltoall: 非固定类型（NBX + 同步发送）
TEST_F(CollectivesTest, AlltoallStringNBXSync)
{
    int rank = comm_->rank();
    int size = comm_->size();
    std::vector<std::string> send_values(size);
    for (int j = 0; j < size; ++j)
        send_values[j] = std::string("from_") + std::to_string(rank) + "_to_" + std::to_string(j);

    auto recv_values =
        alltoall<std::string>(*comm_, send_values, SyncType::NBX, SendMode::Synchronous);
    ASSERT_EQ(static_cast<int>(recv_values.size()), size);
    for (int i = 0; i < size; ++i)
    {
        std::string expected =
            std::string("from_") + std::to_string(i) + "_to_" + std::to_string(rank);
        EXPECT_EQ(recv_values[i], expected);
    }
}